from ship_envs.callbacks.base_callback import BaseCallback
from ship_envs.core.Agent import Agent
from ship_envs.environment.single_ship_env import SingleAgentEnv


def test_callback():
    agent = Agent()
    callback = BaseCallback()
    print(callback.observation_callback(None, agent).shape)


def env_test():
    callback = BaseCallback()
    env = SingleAgentEnv(reset_callback=callback.reset_callback, reward_callback=callback.reward_callback,
                 observation_callback=callback.observation_callback, done_callback=callback.done_callback,
                 info_callback=callback.info_callback)
    # print(env.action_space)
    # print(env.observation_space)
    # print(env.action_space.sample())
    # print(env.observation_space.sample())
    env.reset()
    # env.step(env.action_space.sample())
    # env.step(env.action_space.sample())
    count = 1000
    while count > 0:
        action = env.action_space.sample()
        obs, reward, done, info = env.step(action)
        print(obs.shape)
        env.render()
        count -= 1
        if done:
            print("test done！")
            break


if __name__ == '__main__':
    # test_callback()
    env_test()
